/*
Copyright 2025 PerfXLab (Beijing) Technologies Co., Ltd.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
#include <iostream>
#include <complex>
#include <vector>
#include <cmath>
#include <type_traits>
#include "test_common.hpp"

// ===================== 解决函数歧义的比较工具 =====================
// 每个类型独立的比较函数（彻底解决歧义问题）

// 浮点数比较
inline bool nearly_equal(float a, float b, float eps = 1e-5f) {
    return std::fabs(a - b) < eps;
}

// 双精度浮点数比较
inline bool nearly_equal(double a, double b, double eps = 1e-12) {
    return std::fabs(a - b) < eps;
}

// 单精度复数比较
inline bool nearly_equal(std::complex<float> a, std::complex<float> b, float eps = 1e-5f) {
    return std::fabs(a.real() - b.real()) < eps &&
           std::fabs(a.imag() - b.imag()) < eps;
}

// 双精度复数比较
inline bool nearly_equal(std::complex<double> a, std::complex<double> b, double eps = 1e-12) {
    return std::fabs(a.real() - b.real()) < eps &&
           std::fabs(a.imag() - b.imag()) < eps;
}

// ===================== 类型安全的向量比较函数 =====================
// 根据类型显式分派到正确的比较函数
template <typename T>
bool compare_vectors(size_t n, const T* a, const T* b) {
    // 根据类型选择适当的精度
    constexpr double eps_default = 
        std::is_same_v<T, float> || std::is_same_v<T, std::complex<float>> ? 
        1e-5f : 1e-12;
    
    for (size_t i = 0; i < n; ++i) {
        if constexpr (std::is_same_v<T, float>) {
            if (!nearly_equal(a[i], b[i], static_cast<float>(eps_default)))
                return false;
        }
        else if constexpr (std::is_same_v<T, double>) {
            if (!nearly_equal(a[i], b[i], static_cast<double>(eps_default)))
                return false;
        }
        else if constexpr (std::is_same_v<T, std::complex<float>>) {
            if (!nearly_equal(a[i], b[i], static_cast<float>(eps_default)))
                return false;
        }
        else if constexpr (std::is_same_v<T, std::complex<double>>) {
            if (!nearly_equal(a[i], b[i], static_cast<double>(eps_default)))
                return false;
        }
    }
    return true;
}

// ===================== 向量打印函数 =====================
template <typename T>
void print_vector(const char* name, size_t n, const T* v) {
    std::cout << name << ": [";
    for (size_t i = 0; i < n; ++i) {
        if constexpr (std::is_same_v<T, std::complex<float>> || 
                      std::is_same_v<T, std::complex<double>>) {
            std::cout << "(" << v[i].real() << ", " << v[i].imag() << ")";
        } else {
            std::cout << v[i];
        }
        if (i < n-1) std::cout << ", ";
    }
    std::cout << "]" << std::endl;
}

// ===================== 主测试函数 =====================
TEST(SparseBLASlevel2, scatter) {
    int failures = 0;

    std::cout << "=== Testing OpenSpB_scatter (Fixed Version) ===" << std::endl;

    // ===================== float (sscatter) 测试 =====================
    {
        const int nz = 3;
        float x[] = {1.0f, 2.0f, 3.0f};
        float y_our[5] = {0}; // 初始化为0
        float y_ref[5] = {0}; // 初始化为0
        int indx[] = {0, 2, 4}; // 0-based索引

        // 调用目标函数
        OpenSpB_sscatter(nz, x, y_our, 1, indx, OPENSPBLAS_ZERO_BASE);
        
        // 调用参考函数
        BLAS_sussc(nz, x, y_ref, 1, indx, blas_zero_base);
        
        if (!compare_vectors(5, y_our, y_ref)) {
            std::cout << "FAIL: sscatter test" << std::endl;
            print_vector("Expected y", 5, y_ref);
            print_vector("Actual y", 5, y_our);
            failures++;
        } else {
            std::cout << "PASS: sscatter test" << std::endl;
        }
    }

    // ===================== double (dscatter) 测试 =====================
    {
        const int nz = 2;
        double x[] = {1.5, -2.5};
        double y_our[5] = {1.0, 2.0, 3.0, 4.0, 5.0}; // 有初始值
        double y_ref[5] = {1.0, 2.0, 3.0, 4.0, 5.0}; // 有初始值
        int indx[] = {1, 3}; // 1-based索引

        // 调用目标函数
        OpenSpB_dscatter(nz, x, y_our, 1, indx, OPENSPBLAS_ONE_BASE);
        
        // 调用参考函数
        BLAS_dussc(nz, x, y_ref, 1, indx, blas_one_base);
        
        if (!compare_vectors(5, y_our, y_ref)) {
            std::cout << "FAIL: dscatter test" << std::endl;
            failures++;
        } else {
            std::cout << "PASS: dscatter test" << std::endl;
        }
    }

    // ===================== complex<float> (cscatter) 测试 =====================
    {
        const int nz = 2;
        std::complex<float> x[] = {{1.0f, 2.0f}, {3.0f, -1.0f}};
        std::complex<float> y_our[5] = {{0,0}, {0,0}, {0,0}, {0,0}, {0,0}};
        std::complex<float> y_ref[5] = {{0,0}, {0,0}, {0,0}, {0,0}, {0,0}};
        int indx[] = {0, 3}; // 0-based索引

        // 调用目标函数
        OpenSpB_cscatter(nz, x, y_our, 1, indx, OPENSPBLAS_ZERO_BASE);
        
        // 调用参考函数
        BLAS_cussc(nz, x, y_ref, 1, indx, blas_zero_base);
        
        if (!compare_vectors(5, y_our, y_ref)) {
            std::cout << "FAIL: cscatter test" << std::endl;
            print_vector("Expected y", 5, y_ref);
            print_vector("Actual y", 5, y_our);
            failures++;
        } else {
            std::cout << "PASS: cscatter test" << std::endl;
        }
    }

    // ===================== complex<double> (zscatter) 测试 =====================
    {
        const int nz = 3;
        std::complex<double> x[] = {{1.1, 2.2}, {3.3, 4.4}, {5.5, 6.6}};
        std::complex<double> y_our[4] = {{0,0}, {0,0}, {0,0}, {0,0}};
        std::complex<double> y_ref[4] = {{0,0}, {0,0}, {0,0}, {0,0}};
        int indx[] = {0, 1, 2}; // 0-based索引
        const int incy = 1;      // 使用标准步长

        // 调用目标函数
        OpenSpB_zscatter(nz, x, y_our, incy, indx, OPENSPBLAS_ZERO_BASE);
        
        // 调用参考函数
        BLAS_zussc(nz, x, y_ref, incy, indx, blas_zero_base);
        
        if (!compare_vectors(4, y_our, y_ref)) {
            std::cout << "FAIL: zscatter test" << std::endl;
            print_vector("Expected y", 4, y_ref);
            print_vector("Actual y", 4, y_our);
            failures++;
        } else {
            std::cout << "PASS: zscatter test" << std::endl;
        }
    }

    // ===================== 测试总结 =====================
    std::cout << "\n=== Test Summary ===" << std::endl;
    if (failures == 0) {
        std::cout << "ALL TESTS PASSED!" << std::endl;
    } else {
        std::cout << failures << " TESTS FAILED!" << std::endl;
    }
}